import os
import json
import cv2

from utils.bbox import IOU

label_save_path = "/root/autodl-tmp/sep/figure-separator/benchmark/subFigureImage/labels.json"
predictpath = "/root/autodl-tmp/sep/figure-separator/benchmark/subFigureImage/layout.json"
pictugre_path = "/root/autodl-tmp/sep/figure-separator/benchmark/subFigureImage/images"

class Benchmark:
    def __init__(self, labelpath, predictpath):
        self.labelpath = labelpath
        self.predictpath = predictpath
        self.label = None
        self.predict = dict()
        self.match_result_refer_gt = dict()
        self.loadLabel()
        self.loadPredict()
                
    def loadLabel(self):
        with open(self.labelpath, "r") as fp:
            label = json.load(fp)
        
        self.label = label

                       
    def loadPredict(self):
        fp = open(self.predictpath, "r")
        predict = json.load(fp)
        fp.close()
        pdf_info = predict["pdf_info"]
        predict = dict()
        for idx, page in enumerate(pdf_info):
            preproc_blocks = page["preproc_blocks"]
            page_image = []
            idx = page["page_idx"]
            size = page["page_size"]
            w = size[0]
            h = size[1]
            
            image_path = os.path.join(pictugre_path, str(idx)+".jpg")
            img = cv2.imread(image_path)
            
            o_h, o_w, _ = img.shape
            
            for block in preproc_blocks:
                if block["type"] == "image":
                    try:
                        bbox = block["bbox"]
                        # 还原到原图像尺寸
                        page_image.append([bbox[0]/ w * o_w, bbox[1] / h * o_h, bbox[2]/ w * o_w, bbox[3] / h * o_h])
                    except:
                        print(block)
                        
            
            predict[str(idx)] = page_image
        self.predict = predict
                    
    def run(self, threshold):
        for key in self.predict.keys():
            predicts = self.predict[key]
            labels = self.label[key]
            label_match = self.dealOne(labels, predicts)
            self.match_result_refer_gt[key] = label_match
        
        self.getScore(threshold)    
        
    def getScore(self, threshold):
        # 计算mIOU
        IOU = []
        SFCR = []
        totalIOU = 0
        n = 0
        for key in self.match_result_refer_gt.keys():
            li = self.match_result_refer_gt[key]
            for item in li:
                IOU_item = []
                recgnized = 0
                for match in item:
                    iou = match[-1]
                    IOU_item.append(iou)
                    totalIOU += iou
                    n += 1
                    if iou > threshold:
                        recgnized += 1
                SFCR.append([len(item), recgnized])
                IOU.append(IOU_item)
        
        mIOU = totalIOU / n
        
        totalSFCR = 0         
        for item in SFCR:
            totalSFCR += (item[1] / item[0])
        
        mSFCR = totalSFCR / len(SFCR)
        totalSubfigures, recognizedSubfigures = sum([item[0] for item in SFCR]), sum([item[1] for item in SFCR])
        print(
            f"threshold: {round(threshold, 3)}",
            f"mIOU: {round(mIOU, 3)}",
            f"mSFCR: {round(mSFCR, 3)}",
            f"totalSubfigures: {round(totalSubfigures, 3)}",
            f"recognizedSubfigures: {round(recognizedSubfigures, 3)}"
        )
        return mIOU, mSFCR, totalSubfigures, recognizedSubfigures
                
    def dealOne(self, labels, predicts):
            
        gt_match = []
        for item in labels:
            item_match = []
            for label in item[1]:
                max_iou = IOU(label, predicts[0])
                max_iou_label = labels[0]
                for idx in range(1, len(predicts)):
                    predict = predicts[idx]
                    iou = IOU(predict, label)
                    if iou > max_iou:
                        max_iou = iou
                        max_iou_label = predict
                item_match.append([label, max_iou_label, max_iou])
            gt_match.append(item_match)
        return gt_match
            
if __name__ == "__main__":
    benchmark = Benchmark(label_save_path, predictpath)
    benchmark.run(0.7)
    for i in range(10, 20, 1):
        benchmark.getScore(round(i*0.05, 2))
    
    # with open("./subFIgureImage/layout.json", "r") as fp:
    #     data = json.load(fp)
    